{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/knowledge_graph_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Graph Query Engine\n",
    "\n",
    "Creating a Knowledge Graph usually involves specialized and complex tasks. However, by utilizing the Llama Index (LLM), the KnowledgeGraphIndex, and the GraphStore, we can facilitate the creation of a relatively effective Knowledge Graph from any data source supported by [Llama Hub](https://llamahub.ai/).\n",
    "\n",
    "Furthermore, querying a Knowledge Graph often requires domain-specific knowledge related to the storage system, such as Cypher. But, with the assistance of the LLM and the LlamaIndex KnowledgeGraphQueryEngine, this can be accomplished using Natural Language!\n",
    "\n",
    "In this demonstration, we will guide you through the steps to:\n",
    "\n",
    "- Extract and Set Up a Knowledge Graph using the Llama Index\n",
    "- Query a Knowledge Graph using Cypher\n",
    "- Query a Knowledge Graph using Natural Language"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-wikipedia\n",
    "%pip install llama-index-llms-azure-openai\n",
    "%pip install llama-index-graph-stores-nebula\n",
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-embeddings-azure-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first get ready for basic preparation of Llama Index."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For OpenAI\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "\n",
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(\n",
    "    stream=sys.stdout, level=logging.INFO\n",
    ")  # logging.DEBUG for more verbose output\n",
    "\n",
    "\n",
    "# define LLM\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Azure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "\n",
    "# For Azure OpenAI\n",
    "api_key = \"<api-key>\"\n",
    "azure_endpoint = \"https://<your-resource-name>.openai.azure.com/\"\n",
    "api_version = \"2023-07-01-preview\"\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    model=\"gpt-35-turbo-16k\",\n",
    "    deployment_name=\"my-custom-llm\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=api_version,\n",
    ")\n",
    "\n",
    "# You need to deploy your own embedding model as well as your own chat completion model\n",
    "embed_model = AzureOpenAIEmbedding(\n",
    "    model=\"text-embedding-ada-002\",\n",
    "    deployment_name=\"my-custom-embedding\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=api_version,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare for NebulaGraph\n",
    "\n",
    "Before next step to creating the Knowledge Graph, let's ensure we have a running NebulaGraph with defined data schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ipython-ngql in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (0.5)\n",
      "Requirement already satisfied: nebula3-python in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (3.4.0)\n",
      "Requirement already satisfied: pandas in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython-ngql) (2.0.3)\n",
      "Requirement already satisfied: Jinja2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython-ngql) (3.1.2)\n",
      "Requirement already satisfied: pytz>=2021.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python) (2023.3)\n",
      "Requirement already satisfied: future>=0.18.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python) (0.18.3)\n",
      "Requirement already satisfied: httplib2>=0.20.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python) (0.22.0)\n",
      "Requirement already satisfied: six>=1.16.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python) (1.16.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from httplib2>=0.20.0->nebula3-python) (3.0.9)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from Jinja2->ipython-ngql) (2.1.3)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (2023.3)\n",
      "Requirement already satisfied: numpy>=1.20.3 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (1.25.2)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (2.8.2)\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.2.1 is available.\n",
      "You should consider upgrading via the '/Users/loganmarkewich/llama_index/llama-index/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "# Create a NebulaGraph (version 3.5.0 or newer) cluster with:\n",
    "# Option 0 for machines with Docker: `curl -fsSL nebula-up.siwei.io/install.sh | bash`\n",
    "# Option 1 for Desktop: NebulaGraph Docker Extension https://hub.docker.com/extensions/weygu/nebulagraph-dd-ext\n",
    "\n",
    "# If not, create it with the following commands from NebulaGraph's console:\n",
    "# CREATE SPACE llamaindex(vid_type=FIXED_STRING(256), partition_num=1, replica_factor=1);\n",
    "# :sleep 10;\n",
    "# USE llamaindex;\n",
    "# CREATE TAG entity(name string);\n",
    "# CREATE EDGE relationship(relationship string);\n",
    "# :sleep 10;\n",
    "# CREATE TAG INDEX entity_index ON entity(name(256));\n",
    "\n",
    "%pip install ipython-ngql nebula3-python\n",
    "\n",
    "os.environ[\"NEBULA_USER\"] = \"root\"\n",
    "os.environ[\"NEBULA_PASSWORD\"] = \"nebula\"  # default is \"nebula\"\n",
    "os.environ[\n",
    "    \"NEBULA_ADDRESS\"\n",
    "] = \"127.0.0.1:9669\"  # assumed we have NebulaGraph installed locally\n",
    "\n",
    "space_name = \"llamaindex\"\n",
    "edge_types, rel_prop_names = [\"relationship\"], [\n",
    "    \"relationship\"\n",
    "]  # default, could be omit if create from an empty kg\n",
    "tags = [\"entity\"]  # default, could be omit if create from an empty kg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare for StorageContext with graph_store as NebulaGraphStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import StorageContext\n",
    "from llama_index.graph_stores.nebula import NebulaGraphStore\n",
    "\n",
    "graph_store = NebulaGraphStore(\n",
    "    space_name=space_name,\n",
    "    edge_types=edge_types,\n",
    "    rel_prop_names=rel_prop_names,\n",
    "    tags=tags,\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional)Build the Knowledge Graph with LlamaIndex\n",
    "\n",
    "With the help of Llama Index and LLM defined, we could build Knowledge Graph from given documents.\n",
    "\n",
    "If we have a Knowledge Graph on NebulaGraphStore already, this step could be skipped"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1, load data from Wikipedia for \"Guardians of the Galaxy Vol. 3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import download_loader\n",
    "\n",
    "from llama_index.readers.wikipedia import WikipediaReader\n",
    "\n",
    "loader = WikipediaReader()\n",
    "\n",
    "documents = loader.load_data(\n",
    "    pages=[\"Guardians of the Galaxy Vol. 3\"], auto_suggest=False\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2, Generate a KnowledgeGraphIndex with NebulaGraph as graph_store\n",
    "\n",
    "Then, we will create a KnowledgeGraphIndex to enable Graph based RAG, see [here](https://gpt-index.readthedocs.io/en/latest/examples/index_structs/knowledge_graph/KnowledgeGraphIndex_vs_VectorStoreIndex_vs_CustomIndex_combined.html) for deails, apart from that, we have a Knowledge Graph up and running for other purposes, too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import KnowledgeGraphIndex\n",
    "\n",
    "kg_index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    max_triplets_per_chunk=10,\n",
    "    space_name=space_name,\n",
    "    edge_types=edge_types,\n",
    "    rel_prop_names=rel_prop_names,\n",
    "    tags=tags,\n",
    "    include_embeddings=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a Knowledge Graph on NebulaGraph cluster under space named `llamaindex` about the 'Guardians of the Galaxy Vol. 3' movie, let's play with it a little bit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ipython-ngql in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (0.5)\n",
      "Requirement already satisfied: networkx in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (3.1)\n",
      "Requirement already satisfied: pyvis in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (0.3.2)\n",
      "Requirement already satisfied: Jinja2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython-ngql) (3.1.2)\n",
      "Requirement already satisfied: pandas in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython-ngql) (2.0.3)\n",
      "Requirement already satisfied: nebula3-python in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython-ngql) (3.4.0)\n",
      "Requirement already satisfied: jsonpickle>=1.4.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pyvis) (3.0.1)\n",
      "Requirement already satisfied: ipython>=5.3.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pyvis) (8.10.0)\n",
      "Requirement already satisfied: backcall in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.2.0)\n",
      "Requirement already satisfied: pickleshare in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.7.5)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.30 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (3.0.39)\n",
      "Requirement already satisfied: appnope in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.1.3)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (2.15.1)\n",
      "Requirement already satisfied: traitlets>=5 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (5.9.0)\n",
      "Requirement already satisfied: pexpect>4.3 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (4.8.0)\n",
      "Requirement already satisfied: stack-data in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.6.2)\n",
      "Requirement already satisfied: decorator in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (5.1.1)\n",
      "Requirement already satisfied: jedi>=0.16 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.18.2)\n",
      "Requirement already satisfied: matplotlib-inline in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from ipython>=5.3.0->pyvis) (0.1.6)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from jedi>=0.16->ipython>=5.3.0->pyvis) (0.8.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from Jinja2->ipython-ngql) (2.1.3)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pexpect>4.3->ipython>=5.3.0->pyvis) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from prompt-toolkit<3.1.0,>=3.0.30->ipython>=5.3.0->pyvis) (0.2.6)\n",
      "Requirement already satisfied: six>=1.16.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python->ipython-ngql) (1.16.0)\n",
      "Requirement already satisfied: pytz>=2021.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python->ipython-ngql) (2023.3)\n",
      "Requirement already satisfied: future>=0.18.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python->ipython-ngql) (0.18.3)\n",
      "Requirement already satisfied: httplib2>=0.20.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from nebula3-python->ipython-ngql) (0.22.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.0,!=3.0.1,!=3.0.2,!=3.0.3,<4,>=2.4.2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from httplib2>=0.20.0->nebula3-python->ipython-ngql) (3.0.9)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (2.8.2)\n",
      "Requirement already satisfied: numpy>=1.20.3 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (1.25.2)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from pandas->ipython-ngql) (2023.3)\n",
      "Requirement already satisfied: executing>=1.2.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (1.2.0)\n",
      "Requirement already satisfied: pure-eval in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (0.2.2)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from stack-data->ipython>=5.3.0->pyvis) (2.2.1)\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.2.1 is available.\n",
      "You should consider upgrading via the '/Users/loganmarkewich/llama_index/llama-index/bin/python -m pip install --upgrade pip' command.\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Connection Pool Created\n",
      "INFO:nebula3.logger:Get connection to ('127.0.0.1', 9669)\n",
      "[ERROR]:\n",
      " 'IPythonNGQL' object has no attribute '_decode_value'\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>llamaindex</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         Name\n",
       "0  llamaindex"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# install related packages, password is nebula by default\n",
    "%pip install ipython-ngql networkx pyvis\n",
    "%load_ext ngql\n",
    "%ngql --address 127.0.0.1 --port 9669 --user root --password <password>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nebula3.logger:Get connection to ('127.0.0.1', 9669)\n",
      "INFO:nebula3.logger:Get connection to ('127.0.0.1', 9669)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>e</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>(\"A second trailer for the film\")-[:relationsh...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>(\"Adam McKay\")-[:relationship@-442854342936029...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>(\"Adam McKay\")-[:relationship@8513344855738553...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>(\"Asim Chaudhry\")-[:relationship@-803614038978...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>(\"Bakalova\")-[:relationship@-25325064520311626...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>(\"Bautista\")-[:relationship@-90386029986457371...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>(\"Bautista\")-[:relationship@-90386029986457371...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>(\"Beth Mickle\")-[:relationship@716197657641767...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>(\"Bradley Cooper\")-[:relationship@138630731832...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>(\"Bradley Cooper\")-[:relationship@838402633192...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                   e\n",
       "0  (\"A second trailer for the film\")-[:relationsh...\n",
       "1  (\"Adam McKay\")-[:relationship@-442854342936029...\n",
       "2  (\"Adam McKay\")-[:relationship@8513344855738553...\n",
       "3  (\"Asim Chaudhry\")-[:relationship@-803614038978...\n",
       "4  (\"Bakalova\")-[:relationship@-25325064520311626...\n",
       "5  (\"Bautista\")-[:relationship@-90386029986457371...\n",
       "6  (\"Bautista\")-[:relationship@-90386029986457371...\n",
       "7  (\"Beth Mickle\")-[:relationship@716197657641767...\n",
       "8  (\"Bradley Cooper\")-[:relationship@138630731832...\n",
       "9  (\"Bradley Cooper\")-[:relationship@838402633192..."
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Query some random Relationships with Cypher\n",
    "%ngql USE llamaindex;\n",
    "%ngql MATCH ()-[e]->() RETURN e LIMIT 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nebulagraph_draw.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"500px\"\n",
       "            src=\"nebulagraph_draw.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2b43757f0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# draw the result\n",
    "\n",
    "%ng_draw"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asking the Knowledge Graph\n",
    "\n",
    "Finally, let's demo how to Query Knowledge Graph with Natural language!\n",
    "\n",
    "Here, we will leverage the `KnowledgeGraphQueryEngine`, with `NebulaGraphStore` as the `storage_context.graph_store`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import KnowledgeGraphQueryEngine\n",
    "\n",
    "from llama_index.core import StorageContext\n",
    "from llama_index.graph_stores.nebula import NebulaGraphStore\n",
    "\n",
    "query_engine = KnowledgeGraphQueryEngine(\n",
    "    storage_context=storage_context,\n",
    "    llm=llm,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33;1m\u001b[1;3mGraph Store Query:\n",
      "```\n",
      "MATCH (p:`entity`)-[:relationship]->(m:`entity`) WHERE p.`entity`.`name` == 'Peter Quill'\n",
      "RETURN p.`entity`.`name`;\n",
      "```\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mGraph Store Response:\n",
      "{'p.entity.name': ['Peter Quill', 'Peter Quill', 'Peter Quill', 'Peter Quill', 'Peter Quill']}\n",
      "\u001b[0m\u001b[32;1m\u001b[1;3mFinal Response: \n",
      "\n",
      "Peter Quill is a character in the Marvel Universe. He is the son of Meredith Quill and Ego the Living Planet.\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<b>\n",
       "\n",
       "Peter Quill is a character in the Marvel Universe. He is the son of Meredith Quill and Ego the Living Planet.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Tell me about Peter Quill?\",\n",
    ")\n",
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "```cypher\n",
       "```\n",
       "MATCH (p:`entity`)-[:relationship]->(m:`entity`) \n",
       "  WHERE p.`entity`.`name` == 'Peter Quill'\n",
       "\n",
       "RETURN p.`entity`.`name`;\n",
       "```\n",
       "```\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "graph_query = query_engine.generate_query(\n",
    "    \"Tell me about Peter Quill?\",\n",
    ")\n",
    "\n",
    "graph_query = graph_query.replace(\"WHERE\", \"\\n  WHERE\").replace(\n",
    "    \"RETURN\", \"\\nRETURN\"\n",
    ")\n",
    "\n",
    "display(\n",
    "    Markdown(\n",
    "        f\"\"\"\n",
    "```cypher\n",
    "{graph_query}\n",
    "```\n",
    "\"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could see it helps generate the Graph query:\n",
    "\n",
    "```cypher\n",
    "MATCH (p:`entity`)-[:relationship]->(e:`entity`) \n",
    "  WHERE p.`entity`.`name` == 'Peter Quill' \n",
    "RETURN e.`entity`.`name`;\n",
    "```\n",
    "And synthese the question based on its result:\n",
    "\n",
    "```json\n",
    "{'e2.entity.name': ['grandfather', 'alternate version of Gamora', 'Guardians of the Galaxy']}\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course we still could query it, too! And this query engine could be our best Graph Query Language learning bot, then :)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nebula3.logger:Get connection to ('127.0.0.1', 9669)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>p.entity.name</th>\n",
       "      <th>e.relationship</th>\n",
       "      <th>m.entity.name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Peter Quill</td>\n",
       "      <td>would return to the MCU</td>\n",
       "      <td>May 2021</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Peter Quill</td>\n",
       "      <td>was abducted from Earth</td>\n",
       "      <td>as a child</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Peter Quill</td>\n",
       "      <td>is leader of</td>\n",
       "      <td>Guardians of the Galaxy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Peter Quill</td>\n",
       "      <td>was raised by</td>\n",
       "      <td>a group of alien thieves and smugglers</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Peter Quill</td>\n",
       "      <td>is half-human</td>\n",
       "      <td>half-Celestial</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  p.entity.name           e.relationship  \\\n",
       "0   Peter Quill  would return to the MCU   \n",
       "1   Peter Quill  was abducted from Earth   \n",
       "2   Peter Quill             is leader of   \n",
       "3   Peter Quill            was raised by   \n",
       "4   Peter Quill            is half-human   \n",
       "\n",
       "                            m.entity.name  \n",
       "0                                May 2021  \n",
       "1                              as a child  \n",
       "2                 Guardians of the Galaxy  \n",
       "3  a group of alien thieves and smugglers  \n",
       "4                          half-Celestial  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%ngql \n",
    "MATCH (p:`entity`)-[e:relationship]->(m:`entity`)\n",
    "  WHERE p.`entity`.`name` == 'Peter Quill'\n",
    "RETURN p.`entity`.`name`, e.relationship, m.`entity`.`name`;"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And change the query to be rendered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nebula3.logger:Get connection to ('127.0.0.1', 9669)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>p</th>\n",
       "      <th>e</th>\n",
       "      <th>m</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>(\"Peter Quill\" :entity{name: \"Peter Quill\"})</td>\n",
       "      <td>(\"Peter Quill\")-[:relationship@-84437522554765...</td>\n",
       "      <td>(\"May 2021\" :entity{name: \"May 2021\"})</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>(\"Peter Quill\" :entity{name: \"Peter Quill\"})</td>\n",
       "      <td>(\"Peter Quill\")-[:relationship@-11770408155938...</td>\n",
       "      <td>(\"as a child\" :entity{name: \"as a child\"})</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>(\"Peter Quill\" :entity{name: \"Peter Quill\"})</td>\n",
       "      <td>(\"Peter Quill\")-[:relationship@-79394488349732...</td>\n",
       "      <td>(\"Guardians of the Galaxy\" :entity{name: \"Guar...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>(\"Peter Quill\" :entity{name: \"Peter Quill\"})</td>\n",
       "      <td>(\"Peter Quill\")-[:relationship@325695233021653...</td>\n",
       "      <td>(\"a group of alien thieves and smugglers\" :ent...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>(\"Peter Quill\" :entity{name: \"Peter Quill\"})</td>\n",
       "      <td>(\"Peter Quill\")-[:relationship@555553046209276...</td>\n",
       "      <td>(\"half-Celestial\" :entity{name: \"half-Celestia...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              p  \\\n",
       "0  (\"Peter Quill\" :entity{name: \"Peter Quill\"})   \n",
       "1  (\"Peter Quill\" :entity{name: \"Peter Quill\"})   \n",
       "2  (\"Peter Quill\" :entity{name: \"Peter Quill\"})   \n",
       "3  (\"Peter Quill\" :entity{name: \"Peter Quill\"})   \n",
       "4  (\"Peter Quill\" :entity{name: \"Peter Quill\"})   \n",
       "\n",
       "                                                   e  \\\n",
       "0  (\"Peter Quill\")-[:relationship@-84437522554765...   \n",
       "1  (\"Peter Quill\")-[:relationship@-11770408155938...   \n",
       "2  (\"Peter Quill\")-[:relationship@-79394488349732...   \n",
       "3  (\"Peter Quill\")-[:relationship@325695233021653...   \n",
       "4  (\"Peter Quill\")-[:relationship@555553046209276...   \n",
       "\n",
       "                                                   m  \n",
       "0             (\"May 2021\" :entity{name: \"May 2021\"})  \n",
       "1         (\"as a child\" :entity{name: \"as a child\"})  \n",
       "2  (\"Guardians of the Galaxy\" :entity{name: \"Guar...  \n",
       "3  (\"a group of alien thieves and smugglers\" :ent...  \n",
       "4  (\"half-Celestial\" :entity{name: \"half-Celestia...  "
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%ngql\n",
    "MATCH (p:`entity`)-[e:relationship]->(m:`entity`)\n",
    "  WHERE p.`entity`.`name` == 'Peter Quill'\n",
    "RETURN p, e, m;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nebulagraph_draw.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"500px\"\n",
       "            src=\"nebulagraph_draw.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2ae69fdf0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%ng_draw"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results of this knowledge-fetching query could not be more clear from the renderred graph then."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
